package mq

import (
	"github.com/beanstalkd/go-beanstalk"
	"time"
)

type beanstalkd struct {
	Conn      *beanstalk.Conn
	queueName string        //队列名称
	pri       uint32        //优先级
	delayTime time.Duration //延迟时间
	ttr       time.Duration //等待超时时间
}

func NewBeanstalkd(conn *beanstalk.Conn) *beanstalkd {
	return &beanstalkd{
		Conn: conn,
	}
}

// 设置队列名称
func (p *beanstalkd) QueueName(queueName string) *beanstalkd {
	p.queueName = queueName
	return p
}

// 设置优先级
func (p *beanstalkd) Pri(pri uint32) *beanstalkd {
	p.pri = pri
	return p
}

// 设置延迟时间
func (p *beanstalkd) DelayTime(delayTime time.Duration) *beanstalkd {
	p.delayTime = delayTime
	return p
}

// 设置超时时间
func (p *beanstalkd) Ttr(ttr time.Duration) *beanstalkd {
	p.ttr = ttr
	return p
}

// 队列写入数据
func (p *beanstalkd) Publish(data string) (uint64, error) {
	tube := beanstalk.Tube{Conn: p.Conn, Name: p.queueName} // 选择一个Tube（队列）
	jobData := []byte(data)
	if p.ttr == 0 {
		p.ttr = 10 * time.Second
	}
	id, err := tube.Put(jobData, p.pri, p.delayTime, p.ttr)
	if err != nil {
		return 0, err
	}
	return id, nil
}

// 获取队列执行状态
func (p *beanstalkd) GetStats() (map[string]string, error) {
	data, err := p.Conn.Stats()
	if err != nil {
		return nil, err
	}
	return data, nil
}
